package pos.trigram;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Set;

public class Testing {

	private String mostUsedTag;
	private HashMap<String,HashMap<String,Integer>> tagTagMap;
	private HashMap<String,HashMap<String,Integer>> tagWordMap;
	private HashSet<String> wordSet;
	private HashMap<String, Integer> tagCount;
	private HashMap<String, Integer> tag2Count;

	private static final double LAMBDA = 0.000005;
	private static final double UNK_WORD = 1.0;
	private static final double UNK_TAG = 0.05;

	public Testing(HashMap<String, HashMap<String, Integer>> tagTagMap, HashMap<String, HashMap<String, Integer>> tagWordMap, Set<String> wordSet, HashMap<String,Integer> tagCount, HashMap<String, Integer> tag2Count) {
		this.tagTagMap = tagTagMap;
		this.tagWordMap = tagWordMap;
		this.wordSet = new HashSet<String>();
		Iterator<String> iterator = wordSet.iterator();
		while(iterator.hasNext()) {
			this.wordSet.add(iterator.next());
		}
		this.tagCount = tagCount;
		this.tag2Count = tag2Count;
	}

	/**
	 * @param args
	 */
	public static void main(String[] args) {
		long startTime = System.currentTimeMillis();
		Training training = new Training();
		training.deserialize();
		Testing testing = new Testing(training.getTagTagMap(), training.getTagWordMap(), 
				training.getWordSet(), training.getTagCount(), training.getTag2Count());
		File testFile = null;
		File outputFile = null;
		
		for(int i = 0; i < args.length; i = i+2) {
			if(args[i].equalsIgnoreCase("-t"))
				testFile = new File(args[i+1]);
			if(args[i].equalsIgnoreCase("-o"))
				outputFile = new File(args[i+1]);
		}
		
		testing.test(testFile, outputFile);
		long endTime = System.currentTimeMillis();
		System.out.println((endTime - startTime)/1000 + " seconds");
	}

	private void test(File testFile, File outputFile) {
		try {

			BufferedWriter writer = new BufferedWriter(new FileWriter(outputFile));
			BufferedReader reader = new BufferedReader(new FileReader(testFile));
			String line;

			ArrayList<String> words = new ArrayList<String>(226000);
			while((line = reader.readLine()) != null) {
				words.add(line.trim()/*.toLowerCase()*/);
			}


			Set<String> tagSet = tagTagMap.keySet();
			String[] tagsArray = tagSet.toArray(new String[2]);
			double[][] score = new double[tagsArray.length][words.size()];
			int[][] bptr = new int[tagsArray.length][words.size()];

			for(int t = 0; t < tagsArray.length; t++) {
				score[t][0] = Math.log(probWGT(words.get(0), tagsArray[t]) * probTGT(tagsArray[t], Training.START));
				bptr[t][0] = 0;
			}


			for(int t = 0; t < tagsArray.length; t++) {
				double max = Double.NEGATIVE_INFINITY;
				int bestPrev = -1;
				for(int p = 0; p < tagsArray.length; p++) {
					double probTGT = probTGT(tagsArray[t], tagsArray[p], Training.START);
					double tmpScore = score[p][0] + Math.log(probTGT);
					if(tmpScore > max) {
						max = tmpScore;
						bestPrev = p;
					}
				}
				score[t][1] = max + Math.log(probWGT(words.get(1), tagsArray[t]));
				bptr[t][1] = bestPrev;
			}

			for(int w = 2; w < words.size(); w++) {
				if(w%10000==0) System.out.println(w);
				for(int t = 0; t < tagsArray.length; t++) {
					double max = Double.NEGATIVE_INFINITY;
					int bestPrev = -1;
					for(int p = 0; p < tagsArray.length; p++) {
						int p2 = bptr[p][w-1];
						double probTGT = probTGT(tagsArray[t], tagsArray[p], tagsArray[p2]);
						double tmpScore = score[p][w-1] + Math.log(probTGT);
						if(tmpScore > max) {
							max = tmpScore;
							bestPrev = p;
						}
					}
					score[t][w] = max + Math.log(probWGT(words.get(w), tagsArray[t]));
					bptr[t][w] = bestPrev;
				}
			}

			//			print2D(words, tagsArray, score);
			//			System.out.println(tagWordMap);
			//			System.out.println(tagTagMap);
			//			printStringArray(tagsArray);

			//			print2D(words, tagsArray, bptr);
			//			print2D(words, tagsArray, bptr);

			double max = Double.NEGATIVE_INFINITY;
			int bestEndTag = -1;
			for(int t = 0; t < tagsArray.length; t++) {
				if(score[t][words.size()-1] > max) {
					max = score[t][words.size()-1];
					bestEndTag = t;
				}
			}

			int[] tags = new int[words.size()];
			tags[words.size()-1] = bestEndTag;
			for(int w = words.size()-2; w>=0; w--) {
				tags[w] = bptr[tags[w+1]][w+1];
			}

			for(int w = 0; w < words.size(); w++) {
				String word = words.get(w);
				if(!wordSet.contains(word)) {
					if(Character.isUpperCase(word.charAt(0))) {
						writer.write("NNP" + " " + words.get(w) + "\n");
						continue;
					}
				}
				writer.write(tagsArray[tags[w]] + " " + words.get(w) + "\n");
			}

			writer.flush();
			writer.close();
		} catch (FileNotFoundException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		}
	}

	private void printStringArray(String[] tagsArray) {
		for (int i = 0; i < tagsArray.length; i++) {
			System.out.print(tagsArray[i] + " ");
		}
		System.out.println();
	}

	private void print2D(ArrayList<String> words, String[] tagsArray,
			int[][] bptr) {
		for(int j = 0; j < words.size(); j++)
			System.out.print("\t" + words.get(j));

		for(int i = 0; i < tagsArray.length; i++) {
			System.out.print(tagsArray[i] + "\t");
			for(int j = 0; j < words.size(); j++) {
				System.out.print(bptr[i][j] + "\t");
			}
			System.out.println("");
		}
	}

	private void print2D(ArrayList<String> words, String[] tagsArray,
			double[][] score) {
		for(int j = 0; j < words.size(); j++)
			System.out.print("\t" + words.get(j));
		System.out.println("");

		for(int i = 0; i < tagsArray.length; i++) {
			System.out.print(tagsArray[i] + "\t");
			for(int j = 0; j < words.size(); j++) {
				System.out.print(score[i][j] + "\t");
			}
			System.out.println("");
		}
	}

	/**
	 * probability of a Given Tag wrt to previous Tag
	 * Implemented add 1 smoothing.
	 * @param tag
	 * @param prevTag
	 * @return
	 */
	private double probTGT(String tag, String prevTag) {
		HashMap<String,Integer> hashMap = tagTagMap.get(tag);
		Integer value = hashMap.get(prevTag);
		double ans;
		if(value == null)
			ans = LAMBDA / (tagCount.get(prevTag) + tagCount.size()*LAMBDA);
		else 
			ans = (value + LAMBDA) / (tagCount.get(prevTag) + tagCount.size()*LAMBDA);

		return ans*scaling();
	}

	/**
	 * probability of a Given Tag wrt to previous Tag
	 * Implemented add 1 smoothing.
	 * @param tag
	 * @param prevTag
	 * @return
	 */
	private double probTGT(String tag, String prevTag, String prev2Tag) {
		HashMap<String,Integer> hashMap = tagTagMap.get(tag);
		String prevTagCombined = Training.combineTags(prevTag, prev2Tag);
		Integer value = hashMap.get(prevTagCombined);
		double ans;
		Integer count = tag2Count.get(prevTagCombined);
		if (count!=null) {
			if (value == null)
				ans = LAMBDA / (count + tagCount.size() * LAMBDA);
			else
				ans = (value + LAMBDA) / (count + tagCount.size() * LAMBDA);
		}
		else {
			ans = UNK_TAG; // TODO
		}
		return ans*scaling();
	}

	private double scaling() {
		return 1 /*Math.pow(10,1)*/;
	}

	/**
	 * probability of Word Given Tag
	 * Implemented add 1 smoothing
	 * @param word
	 * @param tag
	 * @return
	 */
	private double probWGT(String word, String tag) {
		HashMap<String,Integer> hashMap = tagWordMap.get(tag);
		Integer value = hashMap.get(word);

		double ans;
		if(wordSet.contains(word)) {
			if(value == null)
				ans = LAMBDA / (tagCount.get(tag) + wordSet.size()*LAMBDA);
			else 
				ans = (value + LAMBDA) / (tagCount.get(tag) + wordSet.size()*LAMBDA);
		}
		else {
			ans = UNK_WORD;
		}

		return ans*scaling();
	}
}
